import os
import numpy as np
import cv2
import matplotlib.pyplot as plt


def draw_curves(save_dir, video_name, x, energies=None, r_errs=None, t_errs=None):
    os.makedirs(os.path.join(save_dir, video_name), exist_ok=True)
    if not energies is None:
        assert len(x) == len(energies)
        plt.clf()
        plt.plot(x, energies, color="red")
        plt.title("Energy")
        plt.xlabel("frame")
        plt.ylabel("")
        plt.savefig(os.path.join(save_dir, video_name, "energy.png"))
    assert len(x) == len(r_errs) and len(x) == len(t_errs)
    plt.clf()
    plt.plot(x, r_errs, color="red", label="r_err (deg)")
    plt.plot(x, t_errs, color="blue", label="t_err (cm)")
    plt.legend()
    plt.title("Error")
    plt.xlabel("frame")
    plt.ylabel("")
    plt.savefig(os.path.join(save_dir, video_name, "error.png"))


def compute_lims_for_pts(all_points):
    corners = np.stack([np.min(all_points, axis=0), np.max(all_points, axis=0)], axis=0)
    center = corners.mean(axis=0)
    max_size = np.max(corners[1] - corners[0]) * 0.5
    lims = np.stack([center - max_size, center + max_size], axis=0).swapaxes(0, 1)
    return lims


def set_axes_equal(ax, limits=None, labels=True):
    '''Make axes of 3D plot have equal scale so that spheres appear as spheres,
    cubes as cubes, etc..  This is one possible solution to Matplotlib's
    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.

    Input
      ax: a matplotlib axis, e.g., as output from plt.gca().
    '''

    ax.set_box_aspect([1, 1, 1])
    if limits is None:
        x_limits = ax.get_xlim3d()
        y_limits = ax.get_ylim3d()
        z_limits = ax.get_zlim3d()

        x_range = abs(x_limits[1] - x_limits[0])
        x_middle = np.mean(x_limits)
        y_range = abs(y_limits[1] - y_limits[0])
        y_middle = np.mean(y_limits)
        z_range = abs(z_limits[1] - z_limits[0])
        z_middle = np.mean(z_limits)

        # The plot bounding box is a sphere in the sense of the infinity
        # norm, hence I call half the max range the plot radius.
        plot_radius = 0.5 * max([x_range, y_range, z_range])
        ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
        ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
        ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])
    else:
        x_limits, y_limits, z_limits = limits
        ax.set_xlim3d([x_limits[0], x_limits[1]])
        ax.set_ylim3d([y_limits[0], y_limits[1]])
        ax.set_zlim3d([z_limits[0], z_limits[1]])

    if labels:
        ax.set_xlabel('X Label')
        ax.set_ylabel('Y Label')
        ax.set_zlabel('Z Label')


def plot_bbox_with_pts(save_path, model_points, pred_points, gt_points, view_angle=[0, 0], save_fig=True, dpi=100, s=2):
    # init
    limits = compute_lims_for_pts(np.concatenate((model_points, pred_points, gt_points), axis=0))
    fig = plt.figure(dpi=dpi, figsize=(6, 6))
    ax = plt.subplot(1, 1, 1, projection='3d', proj_type='ortho')
    ax.view_init(elev=view_angle[0], azim=view_angle[1])
    colors = ['dodgerblue', 'gold', 'red']

    # draw point clouds
    ax.scatter(model_points[:, 0], model_points[:, 1], model_points[:, 2], marker='o', s=s, c=colors[0])
    ax.scatter(pred_points[:, 0], pred_points[:, 1], pred_points[:, 2], marker='o', s=s, c=colors[1])
    ax.scatter(gt_points[:, 0], gt_points[:, 1], gt_points[:, 2], marker='o', s=s, c=colors[2])

    # save figure
    plt.axis('off')
    plt.grid('off')
    set_axes_equal(ax, limits)
    plt.tight_layout()
    fig.savefig(save_path)
    plt.close()


def pcd_bbox_visualization(save_dir, model_points, object_pointss, pred_poses, gt_poses, fps=20):
    os.makedirs(save_dir, exist_ok=True)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    videoWriter = cv2.VideoWriter(os.path.join(save_dir, "tracking_video.mp4"), fourcc=fourcc, fps=fps, frameSize=(600 * 3, 600))
    
    assert len(pred_poses) == len(object_pointss) and len(gt_poses) == len(object_pointss)
    for (object_points, pred_pose, gt_pose) in zip(object_pointss, pred_poses, gt_poses):
        pred_points = (object_points - pred_pose["translation"].reshape(3)) @ pred_pose["rotation"]
        gt_points = (object_points - gt_pose["translation"].reshape(3)) @ gt_pose["rotation"]
        save_path = os.path.join(save_dir, "ex1.png")
        plot_bbox_with_pts(save_path=save_path, model_points=model_points, pred_points=pred_points, gt_points=gt_points, view_angle=[0, 180], save_fig=True)
        img1 = cv2.imread(save_path)
        save_path = os.path.join(save_dir, "ex2.png")
        plot_bbox_with_pts(save_path=save_path, model_points=model_points, pred_points=pred_points, gt_points=gt_points, view_angle=[45, 90], save_fig=True)
        img2 = cv2.imread(save_path)
        save_path = os.path.join(save_dir, "ex3.png")
        plot_bbox_with_pts(save_path=save_path, model_points=model_points, pred_points=pred_points, gt_points=gt_points, view_angle=[90, 0], save_fig=True)
        img3 = cv2.imread(save_path)
        videoWriter.write(np.concatenate((img1, img2, img3), axis=1))
    
    videoWriter.release()
